import math
import os
from datetime import datetime, timedelta
from typing import List, Tuple

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

from utils import get_code_market_data

dir_cache_data = "cache_data"
os.makedirs(dir_cache_data, exist_ok=True)

cached_scaler = (
    {}
)  # 在获取测试集数据时缓存正太分布标准化计算模型，以供测试集标准化时调用


def get_transform_data_history_feature5(
    code_market: str,
    start_date: str,
    end_date: str,
    period: str,
    dividend_type: str,
    pred_window,
    history_window: int,
    is_test: bool,
) -> Tuple[pd.DataFrame, List[str], List[str], List[str]]:
    """
    从指定接口获取股票数据，并拓展截面特征，然后对特征和标签进行正太分布标准化
    :param code_market: 使用的股票代码， 类似"600000.SH"
    :param start: 测试集起始时间，类似 ”20210101“
    :param end: 测试集结束时间，类似 ”20220630“
    :param period: 数据粒度，可以是1天，15分钟，5分钟，1分钟级别，对应输入为 1d, 15m, 5m, 1m
    :param dividend_type: 除权方式，可以是前复权，后复权，等比前复权，等比后复权，对应输入为 front, back, front_ratio, back_ration
    :param pred_window: 标签计算参数，10表示计算十天后股票价格对比明天股票价格的增长率。
    :param history_window: 表示截面模型使用的扩展截面历史价格特征数量，30表示扩展历史30天的价格特征
    :param is_test: 是否是测试集，在数据归一化处理时，测试集根据训练集的数据进行归一化，所以需要区分训练集和测试集
    :return: (特征数据，index列名，特征列名，标签列名)
    """
    modify_start_date = (
        datetime.strptime(start_date, "%Y%m%d")
        - timedelta(days=math.ceil(7 / 5 * history_window + 7))
    ).strftime(
        "%Y%m%d"
    )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露

    modify_end_date = (
        datetime.strptime(end_date, "%Y%m%d")
        + timedelta(days=math.ceil(7 / 5 * pred_window + 2))
    ).strftime(
        "%Y%m%d"
    )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露
    df = get_code_market_data(
        code_market, modify_start_date, modify_end_date, period, dividend_type
    )
    df.sort_values("time", inplace=True)
    df.index.name = "date_time"
    df.reset_index(inplace=True)
    df["date_time"] = df["date_time"].apply(lambda x: str(x))
    index_cols = ["date_time", "time", "code_market"]
    old_feature_cols = ["open", "high", "close", "low", "volume", "price"]
    df["price"] = (df["high"] + df["low"] + df["close"]) / 3
    label_cols = ["change_ratio"]
    df[label_cols[0]] = (
        df["price"].shift(-pred_window - 1) / df["price"].shift(-1)
    ) - 1
    feature_cols = [f"norm_{x}" for x in old_feature_cols]
    if is_test:
        scaler = cached_scaler[code_market]
    else:
        scaler = StandardScaler()
        df_train = df
        scaler.fit(df_train[old_feature_cols + label_cols].dropna())
        cached_scaler[code_market] = scaler

    df.loc[1:, feature_cols + label_cols] = scaler.transform(
        df[old_feature_cols + label_cols]
    )[:-1]
    for i in range(1, history_window):
        df[f"price_{i}"] = df[f"price"].shift(i)
        feature_cols.append(f"price_{i}")
    df.dropna(inplace=True)
    df = df[df["date_time"] <= end_date]
    return df, index_cols, feature_cols, label_cols


cached_scaler2 = {}


def get_transform_data_history_feature6(
    code_market: str,
    start_date: str,
    end_date: str,
    period: str,
    dividend_type: str,
    pred_window,
    history_window: int,
    is_test: bool,
) -> Tuple[pd.DataFrame, List[str], List[str], List[str]]:
    """
    从指定接口获取股票数据，并拓展截面特征，然后对标签进行正太分布标准化
    :param code_market: 使用的股票代码， 类似"600000.SH"
    :param start: 测试集起始时间，类似 ”20210101“
    :param end: 测试集结束时间，类似 ”20220630“
    :param period: 数据粒度，可以是1天，15分钟，5分钟，1分钟级别，对应输入为 1d, 15m, 5m, 1m
    :param dividend_type: 除权方式，可以是前复权，后复权，等比前复权，等比后复权，对应输入为 front, back, front_ratio, back_ration
    :param pred_window: 标签计算参数，10表示计算十天后股票价格对比明天股票价格的增长率。
    :param history_window: 表示截面模型使用的扩展截面历史价格特征数量，30表示扩展历史30天的价格特征
    :param is_test: 是否是测试集，在数据归一化处理时，测试集根据训练集的数据进行归一化，所以需要区分训练集和测试集
    :return: (特征数据，index列名，特征列名，标签列名)
    """
    modify_start_date = (
        datetime.strptime(start_date, "%Y%m%d")
        - timedelta(days=math.ceil(7 / 5 * history_window + 7))
    ).strftime(
        "%Y%m%d"
    )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露

    modify_end_date = (
        datetime.strptime(end_date, "%Y%m%d")
        + timedelta(days=math.ceil(7 / 5 * pred_window + 2))
    ).strftime(
        "%Y%m%d"
    )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露
    df = get_code_market_data(
        code_market, modify_start_date, modify_end_date, period, dividend_type
    )
    df.sort_values("time", inplace=True)
    df.index.name = "date_time"
    df.reset_index(inplace=True)
    df["date_time"] = df["date_time"].apply(lambda x: str(x))
    index_cols = ["date_time", "time", "code_market"]
    old_feature_cols = ["open", "high", "close", "low", "volume", "price"]
    df["price"] = (df["high"] + df["low"] + df["close"]) / 3
    label_cols = ["change_ratio"]
    df[label_cols[0]] = (
        df["price"].shift(-pred_window - 1) / df["price"].shift(-1)
    ) - 1
    feature_cols = [f"norm_{x}" for x in old_feature_cols]
    if is_test:
        scaler = cached_scaler2[code_market]
    else:
        scaler = StandardScaler()

        df_train = df
        scaler.fit(df_train[label_cols].dropna())
        cached_scaler2[code_market] = scaler

    df.loc[1:, label_cols] = scaler.transform(df[label_cols])[:-1]
    df.loc[1:, feature_cols] = df[old_feature_cols].to_numpy()[:-1]
    for i in range(1, history_window):
        df[f"price_{i}"] = df[f"price"].shift(i)
        feature_cols.append(f"price_{i}")
    df.dropna(inplace=True)
    df = df[df["date_time"] <= end_date]
    return df, index_cols, feature_cols, label_cols


feature_deal_map = {
    0: get_transform_data_history_feature6,
    1: get_transform_data_history_feature5,
}


class FeatureDataSource:
    def __init__(
        self,
        code_market_list: List[str],
        start: str,
        end: str,
        period: str = "1m",
        dividend_type: str = "front_ratio",
        pred_window: int = 10,
        history_days: int = 30,
        history_window: int = 30,
        func_id: int = 0,
        is_test: bool = False,
    ):
        """
        :param code_market_list: 使用的股票代码， 类似["600000.SH", .....]
        :param start: 测试集起始时间，类似 ”20210101“
        :param end: 测试集结束时间，类似 ”20220630“
        :param period: 数据粒度，可以是1天，15分钟，5分钟，1分钟级别，对应输入为 1d, 15m, 5m, 1m
        :param dividend_type: 除权方式，可以是前复权，后复权，等比前复权，等比后复权，对应输入为 front, back, front_ratio, back_ration
        :param pred_window: 标签计算参数，10表示计算十天后股票价格对比明天股票价格的增长率。
        :param history_days: 时间序列模型获取特征滚动的天数，10表示滚动10天
        :param history_window: 表示截面模型使用的扩展截面历史价格特征数量，30表示扩展历史30天的价格特征
        :param func_id: 使用的特征处理函数，将原始量价数据转化为特征数据的处理函数，根据一个map映射确定处理的函数
        :param is_test: 是否是测试集，在数据归一化处理时，测试集根据训练集的数据进行归一化，所以需要区分训练集和测试集
        """
        dfs = {}
        self.index_cols, self.feature_cols, self.label_cols = None, None, None
        for code_market in code_market_list:
            # df = get_code_market_data(code_market, "", "", period, dividend_type).reset_index(name="date_time")
            try:
                modify_start_date = (
                    datetime.strptime(start, "%Y%m%d")
                    - timedelta(days=math.ceil(7 / 5 * history_days + 7))
                ).strftime(
                    "%Y%m%d"
                )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露
                df, index_cols, feature_cols, label_cols = feature_deal_map[func_id](
                    code_market,
                    modify_start_date,
                    end,
                    period,
                    dividend_type,
                    pred_window=pred_window,
                    history_window=history_window,
                    is_test=is_test,
                )
                df["code_market"] = code_market
                dfs[code_market] = df.set_index("time").sort_index()
                self.index_cols, self.feature_cols, self.label_cols = (
                    index_cols,
                    feature_cols,
                    label_cols,
                )
            except Exception as e:
                print(f"{code_market} data get errors: {e}")
        all_time_list = []
        for code_market, df in dfs.items():
            all_time_list.extend(df.index.to_list()[history_days - 1 :])
        self.dfs = dfs
        self.all_time_list = sorted(list(set(all_time_list)))  # 去重
        start_timestamp_ms = datetime.strptime(start, "%Y%m%d").timestamp() * 1000
        self.all_time_list = [x for x in self.all_time_list if x > start_timestamp_ms]
        self.start = start
        self.end = end
        self.pred_window = pred_window
        self.history_days = history_days

    def __len__(self):
        """
        :return: 多个股票的一天的数据集合成一个batch，这里返回batch数量
        """
        return len(self.all_time_list)

    def __getitem__(self, item) -> Tuple[np.array, np.array]:
        """
        :param item:
        :return: (input_feature, labels) input_feature是多个股票每一个截面滚动多天的特征，形状为 num_stocks*history_days*feature_size,
        history_days为特征历史滚动天数，适用于时间序列模型。labels为个股票每一个截面的标签，形状为 num_stock

        """
        batch_x = []
        batch_y = []
        cur_time = self.all_time_list[item]
        for xcode, gpdf in self.dfs.items():
            if not cur_time in gpdf.index:
                continue
            last_k = gpdf.index.get_loc(cur_time)
            if last_k < self.history_days - 1:
                continue
            X = (
                gpdf[self.feature_cols]
                .iloc[last_k - self.history_days + 1 : last_k + 1]
                .to_numpy()
            )
            Y = gpdf[self.label_cols].iloc[last_k].to_numpy().squeeze(-1)
            batch_x.append(X)
            batch_y.append(Y)
        batch_x = np.stack(batch_x)
        batch_y = np.stack(batch_y)
        return batch_x, batch_y

    def get_all_items(self) -> Tuple[np.array, np.array, List[List]]:
        """
        :return: (input_feature, labels, a list of time and code_market code),
        input_feature是多个股票每一天截面特征，形状 num_stocks * feature_size
        labels是多个股票这一天截面的标签，我使用的是未来十天股票的增长率，形状 num_stocks
        a list of time and code_market code是每一行截面特征对应的时间戳和股票代码
        """
        df = pd.concat(self.dfs.values())
        X = df[self.feature_cols].to_numpy()
        Y = df[self.label_cols].to_numpy().squeeze(-1)
        df = df.reset_index(names=["time"])
        all_time_code_list = df[["time", "code_market"]].to_numpy().tolist()
        return X, Y, all_time_code_list

    def get_all_item_time_stock(self) -> List[List]:
        """
        :return: a list of time and code_market code, 例如，[[1613664000000, "600000.SH"]....]
        """
        all_time_to_codes = {}
        for cur_time in self.all_time_list:
            for xcode, gpdf in self.dfs.items():
                if not cur_time in gpdf.index:
                    continue
                last_k = gpdf.index.get_loc(cur_time)
                if last_k < self.history_days - 1:
                    continue
                all_time_to_codes.setdefault(cur_time, []).append(xcode)
        time_and_code_list = [[x, y] for x, v in all_time_to_codes.items() for y in v]
        return time_and_code_list

    def get_merge_all_stock(self) -> pd.DataFrame:
        """
        :return: DateFrame with columns ["date_time", "time", "code_market", "open", "close", "high",
        "low", "volume", "openInterest", "change_ratio"]， 示例数据如下：
        20210219,1613664000000,600000.SH,9.01499462807954,9.131531954758316,9.25639337619986,
        8.965050059502921,1229263,15,-0.2537008132351914
        """
        choose_cols = (
            self.index_cols
            + ["open", "close", "high", "low", "volume", "openInterest"]
            + self.label_cols
        )
        data = pd.concat(
            [x.reset_index()[choose_cols] for x in self.dfs.values()], axis=0
        )
        result = data[choose_cols]
        return result
